import numpy as np
import soundfile as sf


def embed_watermark(audio_path, watermark, output_path):
    # 读取并归一化音频（关键修复点）
    audio, sample_rate = sf.read(audio_path)
    # 修改部分：每个声道都嵌入水印
    if audio.ndim > 1:
        for channel in range(audio.shape[1]):
            channel_audio = np.clip(audio[:, channel], -1.0, 1.0)
            # 水印处理
            watermark_bin = ''.join(format(ord(c), '08b') for c in watermark)
            watermark_len = len(watermark_bin)
            if watermark_len > len(channel_audio):
                raise ValueError("水印过长，音频容量不足")

            # 安全转为int16（关键修复点）
            channel_audio_int = np.clip(channel_audio * (2**15), -32768, 32767).astype(np.int16)

            # LSB嵌入
            for i in range(watermark_len):
                if i >= len(channel_audio_int):
                    break
                if(channel_audio_int[i] & 0xFE >= 32767):
                    continue
                channel_audio_int[i] = (channel_audio_int[i] & 0xFE) | int(watermark_bin[i])

            audio[:, channel] = channel_audio_int.astype(np.float32) / 32767
    else:
        print("单声道音频")
    sf.write(output_path, audio.astype(np.float32), sample_rate)

# 测试（当前时间水印）
embed_watermark("original.wav", "2025-03-30-20:24", "output.wav")